import numpy as np
import matplotlib.pyplot as plt
import matplotlib as mpl
from sklearn.linear_model import LinearRegression
from sklearn.metrics import r2_score

# 设置全局字体配置
plt.rcParams['font.size'] = 10
plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei', 'Arial Unicode MS']  # 多字体备选
plt.rcParams['axes.unicode_minus'] = False  # 正确显示负号
plt.rcParams['text.usetex'] = False  # 禁用LaTeX渲染
# 创建支持中文和数学符号的字体属性
chinese_font = mpl.font_manager.FontProperties(fname=mpl.font_manager.findfont('SimHei'))
math_font = mpl.font_manager.FontProperties(fname=mpl.font_manager.findfont('DejaVu Sans'))


# 模拟实验数据（200轮训练）
np.random.seed(42)
rounds = np.arange(1, 201)

# 模拟梯度范数（随训练轮次递减）
gradient_norms = 0.8 * np.exp(-0.015 * rounds) + 0.2 * np.random.normal(0, 0.03, 200)
gradient_norms = np.clip(gradient_norms, 0.1, 1.0)

# 模拟噪声尺度（与梯度范数负相关）
noise_scales = 0.1 + 0.6 * (1 - gradient_norms) + 0.1 * np.random.normal(0, 0.05, 200)
noise_scales = np.clip(noise_scales, 0.1, 0.7)

# 模拟准确率数据
def generate_accuracy(base, noise_scale, start_round=10):
    """生成带噪声的准确率曲线"""
    acc = base * (1 - np.exp(-0.03 * (rounds - start_round)))
    acc[rounds < start_round] = 0
    # 添加与噪声尺度相关的波动
    noise_effect = 0.03 * noise_scale * np.random.normal(0, 1, 200)
    return np.clip(acc + noise_effect, 0, base)

# GP-AdaFL准确率（更高、更稳定）
gp_adafl_acc = generate_accuracy(0.98, noise_scales * 0.5)
gp_adafl_upper = gp_adafl_acc + 0.015 * (1 - noise_scales)
gp_adafl_lower = gp_adafl_acc - 0.015 * (1 - noise_scales)

# DP-FedAvg准确率（较低、波动大）
dp_fedavg_acc = generate_accuracy(0.92, noise_scales * 1.5)
dp_fedavg_upper = dp_fedavg_acc + 0.03 * noise_scales
dp_fedavg_lower = dp_fedavg_acc - 0.03 * noise_scales

# 创建图形
fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(6, 8), dpi=100)

# ================== 子图 (a): 噪声尺度与梯度范数随时间变化曲线 ==================
# 创建双Y轴
ax1a = ax1
ax1b = ax1a.twinx()

# 绘制梯度范数曲线（左轴）
line1, = ax1a.plot(rounds, gradient_norms, 'b-', linewidth=2, label='梯度范数')
ax1a.set_ylabel('梯度范数',  fontsize=11, fontproperties=chinese_font)
ax1a.tick_params(axis='y')
ax1a.set_ylim(0.05, 1.05)

# 绘制噪声尺度曲线（右轴）
line2, = ax1b.plot(rounds, noise_scales, 'r--', linewidth=2, label='噪声尺度σ')
ax1b.set_ylabel('噪声尺度',  fontsize=11, fontproperties=chinese_font)
ax1b.tick_params(axis='y')
ax1b.set_ylim(0.05, 0.75)

# 设置公共属性
ax1a.set_xlabel('训练轮次', fontsize=11, fontproperties=chinese_font)
# ax1a.grid(True, linestyle='--', alpha=0.7)
ax1a.set_xlim(0, 200)

# 添加图例（合并两个轴）
lines = [line1, line2]
labels = [l.get_label() for l in lines]
ax1a.legend(lines, labels, loc='upper center', prop=chinese_font,fontsize=10, frameon=True, framealpha=0.9)

# 添加关键区域标注
ax1a.fill_between([0, 50], 0.05, 1.05, color='green', alpha=0.1, label='训练初期')
ax1a.fill_between([150, 200], 0.05, 1.05, color='purple', alpha=0.1, label='训练后期')
ax1a.text(25, 0.9, '高梯度范数\n低噪声尺度', ha='center', fontproperties=chinese_font, 
          bbox=dict(boxstyle="round,pad=0.3", fc="white", ec="green", alpha=0.8))
ax1a.text(175, 0.3, '低梯度范数\n高噪声尺度', ha='center', fontproperties=chinese_font, 
          bbox=dict(boxstyle="round,pad=0.3", fc="white", ec="purple", alpha=0.8))

# 添加子图标题在下方
ax1a.text(0.5, -0.23, '(a) 噪声尺度与梯度范数随时间变化曲线', 
         fontsize=13, fontweight='bold', 
         ha='center', transform=ax1a.transAxes, 
         fontproperties=chinese_font)

# ================== 子图 (b): 准确率波动带域图 ==================
# 绘制GP-AdaFL准确率曲线和置信区间
ax2.plot(rounds, gp_adafl_acc, 'g-', linewidth=2, label='GP-AdaFL(ours)')
ax2.fill_between(rounds, gp_adafl_lower, gp_adafl_upper, color='green', alpha=0.2)

# 绘制DP-FedAvg准确率曲线和置信区间
ax2.plot(rounds, dp_fedavg_acc, 'm--', linewidth=2, label='DP-FedAvg')
ax2.fill_between(rounds, dp_fedavg_lower, dp_fedavg_upper, color='magenta', alpha=0.2)

# 设置属性
ax2.set_xlabel('训练轮次', fontsize=11, fontproperties=chinese_font)
ax2.set_ylabel('测试准确率', fontsize=11, fontproperties=chinese_font)
ax2.set_ylim(0.5, 1.0)
ax2.set_xlim(0, 200)
# ax2.grid(True, linestyle='--', alpha=0.7)
ax2.legend(loc='lower right', prop=chinese_font, frameon=True, framealpha=0.9)

# 添加关键结论标注
ax2.annotate('GP-AdaFL波动范围: ±1.5%', 
             xy=(120, 0.92), xytext=(50, 0.85),
             arrowprops=dict(arrowstyle='->', color='dimgray'),
             fontsize=10, fontproperties=chinese_font,
             bbox=dict(boxstyle="round,pad=0.3", fc="white", ec="gray", alpha=0.8))

ax2.annotate('DP-FedAvg波动范围: ±3.0%', 
             xy=(80, 0.78), xytext=(30, 0.65),
             arrowprops=dict(arrowstyle='->', color='dimgray'),
             fontsize=10, fontproperties=chinese_font,
             bbox=dict(boxstyle="round,pad=0.3", fc="white", ec="gray", alpha=0.8))

# 添加子图标题在下方
ax2.text(0.5, -0.23, '(b) 准确率波动带域图（GP-AdaFL vs DP-FedAvg）', 
         fontsize=13, fontweight='bold', 
         ha='center', transform=ax2.transAxes, 
         fontproperties=chinese_font)

# ================== 整体设置 ==================
# 添加整体标题
# fig.suptitle('图4：噪声自适应效果分析', 
#              fontsize=16, fontweight='bold', 
#              fontproperties=chinese_font, y=0.95)

# 添加技术标注
fig.text(0.5, 0.05,  # 从0.01改为0.05，向上移动
         f"GP-AdaFL框架 | 数据集: MNIST+CIFAR-10 | 预测窗口大小: 5轮 | R$^2$={r2:.3f}", 
         ha="center", fontsize=9, style='italic', fontproperties=chinese_font)

# 调整布局并保存
plt.tight_layout(rect=[0, 0.05, 1, 0.95])
plt.subplots_adjust(top=0.92, hspace=0.35)  # 增加子图间距
plt.savefig('noise_adaptation_effect.png', bbox_inches='tight', dpi=300)
plt.show()
